import os

import imageio

import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
import glob

import torch
import tqdm
from mpl_toolkits.mplot3d import Axes3D


def draw_head_mask_heatmap(np_data, to_dir=None,
                           mask=None,
                           x_label="head", y_label="layer",
                           title=None):
    sns.set_context({"figure.figsize": (15, 12)})
    sns.set(font_scale=2.5)

    ax = sns.heatmap(np_data, square=True, mask=mask, cmap="RdBu_r",
                     cbar_kws={'orientation': 'vertical', "shrink": 0.8})

    ax.tick_params(labelsize=7)
    ax.set_xticklabels(ax.get_xticklabels(), rotation=20)
    ax.set(title=title)

    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.savefig(to_dir)
    plt.close()


def draw_layer_mask_heatmap(np_data, to_dir=None,
                            mask=None,
                            x_label="head", y_label="layer",
                            title=None):
    dim_0 = np_data.shape[0]
    np_data = np.reshape(np_data, [1, dim_0])

    # sns.set_context({"figure.figsize": (10, 4)})
    sns.set(font_scale=1.8)

    plt.figure(figsize=(10, 4))

    ax = sns.heatmap(np_data, square=True,
                     mask=mask, cmap="RdBu_r",
                     cbar_kws={'orientation': 'horizontal', "shrink": 0.5})

    ax.tick_params(labelsize=7)
    ax.set_xticklabels(ax.get_xticklabels(), rotation=20)
    # ax.set(title="layer weights")

    ax.set(title=title)

    # plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.savefig(to_dir)
    plt.close()


def compose_gif(img_paths, to_dir=None):
    gif_images = []
    for path in tqdm.tqdm(img_paths):
        gif_images.append(imageio.imread(path))
    imageio.mimsave(to_dir, gif_images, fps=3.0)





def draw_head_weights_best(output_dir):
    dir_ = os.path.join(
        output_dir,
        r"head_weights_best_model.bin"
    )

    head_weights = torch.load(
        open(dir_, "rb"),
        map_location=torch.device('cpu')
    )
    head_weights = torch.sigmoid(5e2 * head_weights)
    head_weights = head_weights.detach().cpu().numpy()
    print(head_weights)

    draw_head_mask_heatmap(
        head_weights,
        to_dir=os.path.join(output_dir, "head_weights_best_model.jpg"),
        mask=None,
        x_label="head", y_label="layer",
        title="head weights"
    )


def draw_layer_weights_best(output_dir):
    dir_ = os.path.join(
        output_dir,
        r"layer_weights_best_model.bin"
    )

    layer_weights = torch.load(
        open(dir_, "rb"),
        map_location=torch.device('cpu')
    )
    layer_weights = torch.sigmoid(5e2 * layer_weights)
    layer_weights = layer_weights.detach().cpu().numpy()
    print(layer_weights)

    draw_layer_mask_heatmap(
        layer_weights,
        to_dir=os.path.join(output_dir, "layer_weights_best_model.jpg"),
        mask=None,
        x_label="head", y_label="layer",
        title="layer weights"
    )


def compose_head_weights_gif(output_dir, ):
    list_head_weights_dirs = glob.glob(os.path.join(
        output_dir,
        r"arch_weights/head_weights_step_*.bin"
    ))
    print(list_head_weights_dirs)
    list_head_weights_dirs = sorted(
        list_head_weights_dirs,
        key=lambda x: int(x.split(".")[0].split("_")[-1]),
        reverse=False
    )
    print(list_head_weights_dirs)

    for head_weights_dir_ in list_head_weights_dirs:
        global_step = int(head_weights_dir_.split(".")[0].split("_")[-1])
        head_weights = torch.load(
            open(head_weights_dir_, "rb"),
            map_location=torch.device('cpu')
        )
        head_weights = torch.sigmoid(5e2 * head_weights)
        head_weights = head_weights.detach().cpu().numpy()
        print(head_weights)

        draw_head_mask_heatmap(
            head_weights,
            to_dir=os.path.join(output_dir, "visualization/head_weights_step_%d.jpg" % global_step),
            mask=None,
            x_label="head", y_label="layer",
            title="global_step: %d" % global_step
        )

    # compose gif
    img_paths = glob.glob(os.path.join(
        output_dir,
        r"visualization",
        r"head_weights_step_*.jpg"
    ))
    print(img_paths)
    img_paths = sorted(
        img_paths,
        key=lambda x: int(x.split(".")[0].split("_")[-1]),
        reverse=False
    )
    compose_gif(img_paths, to_dir=os.path.join(output_dir, "head_weights_evolution.gif"))


def compose_layer_weights_gif(output_dir, ):
    list_layer_weights_dirs = glob.glob(os.path.join(
        output_dir,
        r"arch_weights/layer_weights_step_*.bin"
    ))
    print(list_layer_weights_dirs)
    list_layer_weights_dirs = sorted(
        list_layer_weights_dirs,
        key=lambda x: int(x.split(".")[0].split("_")[-1]),
        reverse=False
    )
    print(list_layer_weights_dirs)

    for layer_weights_dir_ in list_layer_weights_dirs:
        global_step = int(layer_weights_dir_.split(".")[0].split("_")[-1])
        layer_weights = torch.load(
            open(layer_weights_dir_, "rb"),
            map_location=torch.device('cpu')
        )
        layer_weights = torch.sigmoid(5e2 * layer_weights)
        layer_weights = layer_weights.detach().cpu().numpy()
        print(layer_weights)

        draw_layer_mask_heatmap(
            layer_weights,
            to_dir=os.path.join(output_dir, "visualization/layer_weights_step_%d.jpg" % global_step),
            mask=None,
            x_label="head", y_label="layer",
            title="global_step: %d" % global_step
        )

    # compose gif
    img_paths = glob.glob(os.path.join(
        output_dir,
        r"visualization",
        r"layer_weights_step_*.jpg"
    ))
    print(img_paths)
    img_paths = sorted(
        img_paths,
        key=lambda x: int(x.split(".")[0].split("_")[-1]),
        reverse=False
    )
    compose_gif(img_paths, to_dir=os.path.join(output_dir, "layer_weights_evolution.gif"))



##############
# 3d柱状图
##############

def bar3d_graph(np_data, to_dir=None):
    fig = plt.figure(dpi=100)
    ax = fig.add_subplot(111, projection='3d')
    color_list = plt.cm.Blues(np.linspace(0.5, 1.3, 12))
    for z, c in zip(range(len(np_data)), color_list):
        ax.bar3d(range(12), z, np.zeros(12),
                 dx=0.5, dy=0.5, dz=np_data[z], alpha=0.7, color=c, shade=False)
    ax.set_xlabel("head")
    ax.set_ylabel("layer")
    ax.set_zlabel("variance")

    ax.set_zlim(zmin=0.0, zmax=0.01)

    ax.set(title=" head weights' variances")
    # plt.show()

    plt.savefig(to_dir)
    plt.close()


if __name__ == "__main__":

    # img_paths = ["tmp/peggy.jpeg", "tmp/peggy_1.jpeg", ]
    # compose_gif(img_paths)

    # data = np.random.rand(12, 12)
    # mask = np.array([[1, 1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 1] for _ in range(12)], dtype=bool)
    #
    # draw_head_mask_heatmap(data, to_dir="tmp/heatmap.jpg",
    #                        mask=mask,
    #                        x_label="head", y_label="layer")
    #
    # data = np.random.rand(1, 12)
    # print(data.shape)
    #
    # draw_layer_mask_heatmap(data,
    #                         to_dir="tmp/heatmap_0.jpg",
    #                         mask=None,
    #                         x_label="head", y_label="layer")

    # output_dir = r"experiments/outputs/mrpc_32500"
    # # now head mask gif:
    # compose_head_weights_gif(output_dir, )
    # # 生成 layer weights gif
    # compose_layer_weights_gif(output_dir, )
    # # head weights of best model
    # draw_head_weights_best(output_dir)
    # # layer weights of best model
    # draw_layer_weights_best(output_dir)

    # output_dir = r"experiments/outputs/sent_pair_32500"
    # # now head mask gif:
    # compose_head_weights_gif(output_dir, )
    # # 生成 layer weights gif
    # compose_layer_weights_gif(output_dir, )
    # # head weights of best model
    # draw_head_weights_best(output_dir)
    # # layer weights of best model
    # draw_layer_weights_best(output_dir)


    # # head weights from the run_bertology
    #
    # head_importance = np.load("tmp/head_importance.npy")
    # print(head_importance)
    # print(np.sum(head_importance, axis=0))
    # print(np.sum(head_importance, axis=1))
    #
    # draw_head_mask_heatmap(head_importance, to_dir="tmp/head_importance_sixteen_heads_rte.jpg",
    #                        mask=None,
    #                        x_label="head", y_label="layer",
    #                        title="head weights")


    # 画3d柱状图
    # steps = [200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100]
    steps = [i * 50 + 200 for i in range(1, 15)]

    list_head_weights = []
    for step_ in steps:
        dir_ = os.path.join(
            "experiments/outputs/sent_pair_32500",
            "arch_weights/head_weights_step_%d.bin" % step_
        )
        head_weights = torch.load(
            open(dir_, "rb"),
            map_location=torch.device('cpu')
        )
        head_weights = torch.sigmoid(5e2 * head_weights)
        head_weights = head_weights.detach().cpu().numpy()
        # print(head_weights)
        list_head_weights.append(head_weights)

    list_head_weights = np.array(list_head_weights)
    print(list_head_weights.shape)

    head_weights_variances = np.var(list_head_weights, axis=0, keepdims=False)
    print(head_weights_variances.shape)

    head_weights_variances = head_weights_variances / 4

    # to_dir = "tmp/variances_bi_level.jpg"
    to_dir = "tmp/variances_clo.jpg"
    bar3d_graph(head_weights_variances, to_dir=to_dir)




